import vizdoom as vzd
import random
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from skimage import transform# Help us to preprocess the frames

from collections import deque# Ordered collection with ends
import matplotlib.pyplot as plt # Display graphs

import warnings # This ignore all the warning messages that are normally printed during the training because of skiimage
warnings.filterwarnings('ignore')


### MODEL HYPERPARAMETERS
state_size = [84,84,4]      # Our input is a stack of 4 frames hence 84x84x4 (Width, height, channels)
# action_size = game.get_available_buttons_size()              # 3 possible actions: left, right, shoot
action_size = 3
learning_rate =  0.0002      # Alpha (aka learning rate)

### TRAINING HYPERPARAMETERS
total_episodes = 1000        # Total episodes for training
max_steps = 300              # Max possible steps in an episode
batch_size = 64

# Exploration parameters for epsilon greedy strategy
explore_start = 0.5            # exploration probability at start
explore_stop = 0.01            # minimum exploration probability
decay_rate = 0.001            # exponential decay rate for exploration prob

# Q learning hyperparameters
gamma = 0.95               # Discounting rate

### MEMORY HYPERPARAMETERS
pretrain_length = batch_size   # Number of experiences stored in the Memory when initialized for the first time
memory_size = 100000          # Number of experiences the Memory can keep

### MODIFY THIS TO FALSE IF YOU JUST WANT TO SEE THE TRAINED AGENT
training = True

## TURN THIS TO TRUE IF YOU WANT TO RENDER THE ENVIRONMENT
episode_render = False


"""
Here we create our environment
"""


def create_environment():
    game = vzd.DoomGame()

    # Load the correct configuration
    game.load_config("./doom_cfg/basic.cfg")
    # Load the correct scenario (in our case basic scenario)
    game.set_doom_scenario_path("./doom_cfg/basic.wad")

    game.init()

    # Here our possible actions
    left = [1, 0, 0]
    right = [0, 1, 0]
    shoot = [0, 0, 1]
    possible_actions = [left, right, shoot]

    return game, possible_actions


"""
Here we performing random action to test the environment
"""


def test_environment():
    game = vzd.DoomGame()
    game.load_config("./doom_cfg/basic.cfg")
    game.set_doom_scenario_path("./doom_cfg/basic.wad")
    game.init()
    shoot = [0, 0, 1]
    left = [1, 0, 0]
    right = [0, 1, 0]
    actions = [shoot, left, right]

    episodes = 10
    for i in range(episodes):
        game.new_episode()
        while not game.is_episode_finished():
            state = game.get_state()
            img = state.screen_buffer
            misc = state.game_variables
            action = random.choice(actions)
            print(action)
            reward = game.make_action(action)
            print("\treward:", reward)
            time.sleep(0.02)
        print("Result:", game.get_total_reward())
        time.sleep(2)
    game.close()


"""
    preprocess_frame:
    Take a frame.
    Resize it.
        __________________
        |                 |
        |                 |
        |                 |
        |                 |
        |_________________|

        to
        _____________
        |            |
        |            |
        |            |
        |____________|
    Normalize it.

    return preprocessed_frame

    """


def preprocess_frame(frame):
    # Greyscale frame already done in our vizdoom config
    # x = np.mean(frame,-1)
    frame = np.mean(frame,axis=0).astype(np.float32)

    # Crop the screen (remove the roof because it contains no information)
    cropped_frame = frame[30:-10, 30:-30]

    # Normalize Pixel Values
    normalized_frame = cropped_frame / 255.0

    # Resize
    preprocessed_frame = transform.resize(normalized_frame, [84, 84])

    return preprocessed_frame


stack_size = 4  # We stack 4 frames



def stack_frames(stacked_frames, state, is_new_episode):
    # Preprocess frame
    frame = preprocess_frame(state)

    if is_new_episode:
        # Clear our stacked_frames
        stacked_frames = deque([np.zeros((3, 84, 84), dtype=np.int) for i in range(stack_size)], maxlen=4)

        # Because we're in a new episode, copy the same frame 4x
        stacked_frames.append(frame)
        stacked_frames.append(frame)
        stacked_frames.append(frame)
        stacked_frames.append(frame)

        # Stack the frames
        stacked_state = np.stack(stacked_frames, axis=0)

    else:
        # Append frame to deque, automatically removes the oldest frame
        stacked_frames.append(frame)

        # Build the stacked state (first dimension specifies different frames)
        stacked_state = np.stack(stacked_frames, axis=0)

    return torch.from_numpy(stacked_state), stacked_frames


class DQN(nn.Module):
    def __init__(self):
        super(DQN,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=4,out_channels=32,kernel_size=8,stride=4)
        self.bn1 = nn.BatchNorm2d(32)

        self.conv2 = nn.Conv2d(in_channels=32,out_channels=64,kernel_size=4,stride=2)
        self.bn2 = nn.BatchNorm2d(64)

        self.conv3 = nn.Conv2d(in_channels=64,out_channels=128,kernel_size=4,stride=2)
        self.bn3 = nn.BatchNorm2d(128)

        self.fc1 = nn.Linear(3 * 3 * 128,512)
        self.fc2 = nn.Linear(512,3)

    def forward(self, x):
        x = F.elu(self.bn1(self.conv1(x)))
        x = F.elu(self.bn2(self.conv2(x)))
        x = F.elu(self.bn3(self.conv3(x)))
        x = x.view(-1,3 * 3 * 128)
        x = self.fc2(self.fc1(x))
        return x


class Memory():
    def __init__(self, max_size):
        self.buffer = deque(maxlen=max_size)

    def add(self, experience):
        self.buffer.append(experience)

    def sample(self, batch_size):
        buffer_size = len(self.buffer)
        index = np.random.choice(np.arange(buffer_size),
                                 size=batch_size,
                                 replace=False)

        return [self.buffer[i] for i in index]


game, possible_actions = create_environment()

# Instantiate memory
memory = Memory(max_size=memory_size)

model = DQN()

model.load_state_dict(torch.load('./save/model.pkl'))

loss_func = nn.MSELoss()

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

"""
This function will do the part
With ϵ select a random action atat, otherwise select at=argmaxaQ(st,a)
"""


def predict_action(explore_start, explore_stop, decay_rate, decay_step, state, actions):
    ## EPSILON GREEDY STRATEGY
    # Choose action a from state s using epsilon greedy.
    ## First we randomize a number
    exp_exp_tradeoff = np.random.rand()

    # Here we'll use an improved version of our epsilon greedy strategy used in Q-learning notebook
    explore_probability = explore_stop + (explore_start - explore_stop) * np.exp(-decay_rate * decay_step)

    if (explore_probability > exp_exp_tradeoff):
        # Make a random action (exploration)
        action = random.choice(possible_actions)

    else:
        # Get action from Q-network (exploitation)
        # Estimate the Qs values state
        Qs = model(state.unsqueeze(0))

        # Take the biggest Q value (= the best action)
        choice = torch.argmax(Qs)
        action = possible_actions[int(choice)]

    return action, explore_probability

def pretrain(stacked_frames):
    # Instantiate memory

    # Render the environment
    game.new_episode()

    for i in range(pretrain_length):
        # If it's the first step
        if i == 0:
            # First we need a state
            state = game.get_state().screen_buffer
            state, stacked_frames = stack_frames(stacked_frames, state, True)

        # Random action
        action = random.choice(possible_actions)

        # Get the rewards
        reward = game.make_action(action)

        # Look if the episode is finished
        done = game.is_episode_finished()

        # If we're dead
        if done:
            # We finished the episode
            next_state = torch.zeros(state.shape)

            # Add experience to memory
            memory.add((state, action, reward, next_state, done))

            # Start a new episode
            game.new_episode()

            # First we need a state
            state = game.get_state().screen_buffer

            # Stack the frames
            state, stacked_frames = stack_frames(stacked_frames, state, True)
        else:
            # Get the next state
            next_state = game.get_state().screen_buffer
            next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)

            # Add experience to memory
            memory.add((state, action, reward, next_state, done))

            # Our state is now the next_state
            state = next_state


def train(stacked_frames):
    decay_step = 0
    game.init()

    for episode in range(total_episodes):
        # Set step to 0
        step = 0

        # Initialize the rewards of the episode
        episode_rewards = []

        # Make a new episode and observe the first state
        game.new_episode()
        state = game.get_state().screen_buffer

        # Remember that stack frame function also call our preprocess function.
        state, stacked_frames = stack_frames(stacked_frames, state, True)

        while step < max_steps:
            step += 1

            # Increase decay_step
            decay_step += 1

            # Predict the action to take and take it
            action, explore_probability = predict_action(explore_start, explore_stop, decay_rate, decay_step, state,
                                                         possible_actions)

            # Do the action
            reward = game.make_action(action)

            # Look if the episode is finished
            done = game.is_episode_finished()

            # Add the reward to total reward
            episode_rewards.append(reward)

            # If the game is finished
            if done:
                # the episode ends so no next state
                next_state = np.zeros((3, 84, 84), dtype=np.int)
                next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)

                # Set step = max_steps to end the episode
                step = max_steps

                # Get the total reward of the episode
                total_reward = np.sum(episode_rewards)

                print('Episode: {}'.format(episode),
                      'Total reward: {}'.format(total_reward),
                      'Training loss: {:.4f}'.format(loss),
                      'Explore P: {:.4f}'.format(explore_probability))

                memory.add((state, action, reward, next_state, done))

            else:
                # Get the next state
                next_state = game.get_state().screen_buffer

                # Stack the frame of the next_state
                next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)

                # Add experience to memory
                memory.add((state, action, reward, next_state, done))

                # st+1 is now our current state
                state = next_state

            ### LEARNING PART
            # Obtain random mini-batch from memory
            batch = memory.sample(batch_size)
            states_mb = torch.cat([each[0].unsqueeze(0) for each in batch],dim=0)
            actions_mb = torch.cat([torch.as_tensor(each[1]).unsqueeze(0) for each in batch],dim=0)
            rewards_mb = torch.cat([torch.as_tensor(each[2]).unsqueeze(0) for each in batch],dim=0).view(-1,1)
            next_states_mb = torch.cat([each[3].unsqueeze(0) for each in batch],dim=0)
            dones_mb = torch.cat([torch.as_tensor(each[4]).unsqueeze(0) for each in batch],dim=0).view(-1,1)

            states_mb = Variable(states_mb)
            actions_mb = Variable(actions_mb)
            rewards_mb = Variable(rewards_mb)
            dones_mb = Variable(dones_mb)
            next_states_mb = Variable(next_states_mb)

            target_Qs_batch = []

            # Get Q values for next_state
            Qs_next_state = model(next_states_mb)

            # Set Q_target = r if the episode ends at s+1, otherwise set Q_target = r + gamma*maxQ(s', a')
            for i in range(0, len(batch)):
                terminal = dones_mb[i]

                # If we are in a terminal state, only equals reward
                if terminal:
                    target_Qs_batch.append(rewards_mb[i])

                else:
                    target = rewards_mb[i] + gamma * torch.max(Qs_next_state[i])
                    target_Qs_batch.append(target)

            targets_mb = Variable(torch.as_tensor(target_Qs_batch).view(batch_size,1))
            old_targets_mb = torch.bmm(model(states_mb).view(batch_size, 1, -1), actions_mb.view(batch_size, -1, 1).float()).view(batch_size, 1)

            loss = loss_func(targets_mb,old_targets_mb)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # Save model every 5 episodes
        if episode % 100 == 0:
            torch.save(model.state_dict(),'./save/model.pkl')
            print("Model Saved")



if __name__ == '__main__':
    # test_environment()
    # test_input = torch.randn((8,4,84,84))
    # dqn = DQN()
    # test_output = dqn(test_input)
    # print(test_output.shape)

    # Initialize deque with zero-images one array for each image
    stacked_frames = deque([np.zeros((84, 84), dtype=np.int) for i in range(stack_size)], maxlen=4)
    pretrain(stacked_frames)
    train(stacked_frames)